import fs from "fs";
import path from "path";
import { OllamaEmbeddings } from "@langchain/ollama";
import { DocxLoader } from "@langchain/community/document_loaders/fs/docx";

import { MemoryVectorStore } from "langchain/vectorstores/memory";
import { RecursiveCharacterTextSplitter } from "@langchain/textsplitters";

import { pull } from "langchain/hub";
import { Annotation } from "@langchain/langgraph";

export default class EmbeddingsModal {
  memoStorage;
  embedding;
  textSplitter;
  annotation;

  /**
   * 获取 rag 回复模板
   * **/
  async getRagProgramTemplate() {
    return await pull("rlm/rag-prompt");
  }

  constructor() {
    this.embedding = new OllamaEmbeddings({
      model: "nomic-embed-text:latest",
      baseUrl: process.env.MODEL_URL,
    });
    this.memoStorage = new MemoryVectorStore(this.embedding);
    this.textSplitter = new RecursiveCharacterTextSplitter({
      chunkSize: 1000,
      chunkOverlap: 200,
    });
    this.annotation = Annotation.Root({
      question: Annotation,
      context: Annotation,
      answer: Annotation,
    });
  }

  /**
   * 加载本地文件
   * **/
  loadFile(dirPath) {
    console.log("load local file start...");
    if (!fs.existsSync(dirPath))
      throw new Error("search file path can not be exists...");
    const loader = new DocxLoader(dirPath);
    console.log("load local file end...");
    return loader.load();
  }

  /**
   * 文档拆分，内容转向量数据库
   * **/
  async docsSplitter(fileContent) {
    console.log("split text file start...");
    return await this.textSplitter.splitDocuments(fileContent);
  }

  /**
   * 存储向量数据库
   * **/
  async storageStorage(value) {
    console.log("storage db start...");
    await this.memoStorage.addDocuments(value);
    console.log("storage db end...");
  }

  /**
   * 根据问题查询上下文，向量数据库检索
   * **/
  async searchContext(state) {
    console.log("search db start...");
    const startDate = Date.now();
    const retrievedDocs = await this.memoStorage.similaritySearch(
      state.question
    );
    console.log("search db end, exec time is:", Date.now() - startDate);
    return { context: retrievedDocs };
  }

  async startRagSearch() {
    try {
      console.log("ready search flow start...");
      const fileContent = await this.loadFile(
        path.resolve("C:/Users/AO237/Desktop/search/search.docx")
      );
      const splitterValue = await this.docsSplitter(fileContent);
      await this.storageStorage(splitterValue);
      console.log("ready search flow end...");
    } catch (e) {
      console.log("ready search flow error...", e);
    }
  }
}
